# !/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time        : 2021/10/21 19:07
@Author      : Albert Darren
@Contact     : 2563491540@qq.com
@File        : ID3_decision_tree.py
@Version     : Version 1.0.0
@Description : TODO
@Created By  : PyCharm
"""

from numpy import *
import numpy as np
import pandas as pd
from math import log
import operator


# 计算数据集的香农熵
def calc_shannon_ent(data_set):
    num_entries = len(data_set)
    label_counts = {}
    # 给所有可能分类创建字典
    for featVec in data_set:
        current_label = featVec[-1]
        if current_label not in label_counts.keys():
            label_counts[current_label] = 0
        label_counts[current_label] += 1
    shannon_ent = 0.0
    # 以2为底数计算香农熵
    for key in label_counts:
        prob = float(label_counts[key]) / num_entries
        shannon_ent -= prob * log(prob, 2)
    return shannon_ent


# 对离散变量划分数据集，取出该特征取值为value的所有样本
def split_data_set(data_set, axis, value):
    ret_data_set = []
    for featVec in data_set:
        if featVec[axis] == value:
            reduced_feat_vec = featVec[:axis]
            reduced_feat_vec.extend(featVec[axis + 1:])
            ret_data_set.append(reduced_feat_vec)
    return ret_data_set


# 对连续变量划分数据集，direction规定划分的方向，
# 决定是划分出小于value的数据样本还是大于value的数据样本集
def split_continuous_data_set(data_set, axis, value, direction):
    ret_data_set = []
    for featVec in data_set:
        if direction == 0:
            if featVec[axis] > value:
                reduced_feat_vec = featVec[:axis]
                reduced_feat_vec.extend(featVec[axis + 1:])
                ret_data_set.append(reduced_feat_vec)
        else:
            if featVec[axis] <= value:
                reduced_feat_vec = featVec[:axis]
                reduced_feat_vec.extend(featVec[axis + 1:])
                ret_data_set.append(reduced_feat_vec)
    return ret_data_set


# 选择最好的数据集划分方式
def choose_best_feature_to_split(data_set, labels):
    num_features = len(data_set[0]) - 1
    base_entropy = calc_shannon_ent(data_set)
    best_info_gain = 0.0
    best_feature = -1
    best_split_dict = {}
    for i in range(num_features):
        feat_list = [example[i] for example in data_set]
        # 对连续型特征进行处理
        if type(feat_list[0]).__name__ == 'float' or type(feat_list[0]).__name__ == 'int':
            # 产生n-1个候选划分点
            sort_feat_list = sorted(feat_list)
            split_list = []
            for j in range(len(sort_feat_list) - 1):
                split_list.append((sort_feat_list[j] + sort_feat_list[j + 1]) / 2.0)

            best_split_entropy = 10000
            slen = len(split_list)
            # 求用第j个候选划分点划分时，得到的信息熵，并记录最佳划分点
            for j in range(slen):
                value = split_list[j]
                newEntropy = 0.0
                subDataSet0 = split_continuous_data_set(data_set, i, value, 0)
                subDataSet1 = split_continuous_data_set(data_set, i, value, 1)
                prob0 = len(subDataSet0) / float(len(data_set))
                newEntropy += prob0 * calc_shannon_ent(subDataSet0)
                prob1 = len(subDataSet1) / float(len(data_set))
                newEntropy += prob1 * calc_shannon_ent(subDataSet1)
                if newEntropy < best_split_entropy:
                    best_split_entropy = newEntropy
                    bestSplit = j
            # 用字典记录当前特征的最佳划分点
            best_split_dict[labels[i]] = split_list[bestSplit]
            infoGain = base_entropy - best_split_entropy
        # 对离散型特征进行处理
        else:
            uniqueVals = set(feat_list)
            newEntropy = 0.0
            # 计算该特征下每种划分的信息熵
            for value in uniqueVals:
                subDataSet = split_data_set(data_set, i, value)
                prob = len(subDataSet) / float(len(data_set))
                newEntropy += prob * calc_shannon_ent(subDataSet)
            infoGain = base_entropy - newEntropy
        if infoGain > best_info_gain:
            best_info_gain = infoGain
            best_feature = i
    # 若当前节点的最佳划分特征为连续特征，则将其以之前记录的划分点为界进行二值化处理
    # 即是否小于等于bestSplitValue
    if type(data_set[0][best_feature]).__name__ == 'float' or type(data_set[0][best_feature]).__name__ == 'int':
        bestSplitValue = best_split_dict[labels[best_feature]]
        labels[best_feature] = labels[best_feature] + '<=' + str(bestSplitValue)
        for i in range(shape(data_set)[0]):
            if data_set[i][best_feature] <= bestSplitValue:
                data_set[i][best_feature] = 1
            else:
                data_set[i][best_feature] = 0
    return best_feature


# 特征若已经划分完，节点下的样本还没有统一取值，则需要进行投票
def majorityCnt(classList):
    classCount = {}
    for vote in classList:
        if vote not in classCount.keys():
            classCount[vote] = 0
        classCount[vote] += 1
    return max(classCount)


# 主程序，递归产生决策树
def createTree(dataSet, labels, data_full, labels_full):
    classList = [example[-1] for example in dataSet]
    if classList.count(classList[0]) == len(classList):
        return classList[0]
    if len(dataSet[0]) == 1:
        return majorityCnt(classList)
    bestFeat = choose_best_feature_to_split(dataSet, labels)
    bestFeatLabel = labels[bestFeat]
    myTree = {bestFeatLabel: {}}
    featValues = [example[bestFeat] for example in dataSet]
    uniqueVals = set(featValues)
    if type(dataSet[0][bestFeat]).__name__ == 'str':
        currentlabel = labels_full.index(labels[bestFeat])
        featValuesFull = [example[currentlabel] for example in data_full]
        uniqueValsFull = set(featValuesFull)
    del (labels[bestFeat])
    # 针对bestFeat的每个取值，划分出一个子树。
    for value in uniqueVals:
        subLabels = labels[:]
        if type(dataSet[0][bestFeat]).__name__ == 'str':
            uniqueValsFull.remove(value)
        myTree[bestFeatLabel][value] = createTree(split_data_set \
                                                      (dataSet, bestFeat, value), subLabels, data_full, labels_full)
    if type(dataSet[0][bestFeat]).__name__ == 'str':
        for value in uniqueValsFull:
            myTree[bestFeatLabel][value] = majorityCnt(classList)
    return myTree


import matplotlib.pyplot as plt

decisionNode = dict(boxstyle="sawtooth", fc="0.8")
leafNode = dict(boxstyle="round4", fc="0.8")
arrow_args = dict(arrowstyle="<-")


# 计算树的叶子节点数量
def getNumLeafs(myTree):
    numLeafs = 0
    firstStr = myTree.keys()[0]
    secondDict = myTree[firstStr]
    for key in secondDict.keys():
        if type(secondDict[key]).__name__ == 'dict':
            numLeafs += getNumLeafs(secondDict[key])
        else:
            numLeafs += 1
    return numLeafs


# 计算树的最大深度
def getTreeDepth(myTree):
    maxDepth = 0
    firstStr = myTree.keys()[0]
    secondDict = myTree[firstStr]
    for key in secondDict.keys():
        if type(secondDict[key]).__name__ == 'dict':
            thisDepth = 1 + getTreeDepth(secondDict[key])
        else:
            thisDepth = 1
        if thisDepth > maxDepth:
            maxDepth = thisDepth
    return maxDepth


# 画节点
def plotNode(nodeTxt, centerPt, parentPt, nodeType):
    createPlot.ax1.annotate(nodeTxt, xy=parentPt, xycoords='axes fraction', xytext=centerPt, textcoords='axes fraction',
                            va="center", ha="center",
                            bbox=nodeType, arrowprops=arrow_args)


# 画箭头上的文字
def plotMidText(cntrPt, parentPt, txtString):
    lens = len(txtString)
    xMid = (parentPt[0] + cntrPt[0]) / 2.0 - lens * 0.002
    yMid = (parentPt[1] + cntrPt[1]) / 2.0
    createPlot.ax1.text(xMid, yMid, txtString)


def plotTree(myTree, parentPt, nodeTxt):
    numLeafs = getNumLeafs(myTree)
    depth = getTreeDepth(myTree)
    firstStr = myTree.keys()[0]
    cntrPt = (plotTree.x0ff + (1.0 + float(numLeafs)) / 2.0 / plotTree.totalW, plotTree.y0ff)
    plotMidText(cntrPt, parentPt, nodeTxt)
    plotNode(firstStr, cntrPt, parentPt, decisionNode)
    secondDict = myTree[firstStr]
    plotTree.y0ff = plotTree.y0ff - 1.0 / plotTree.totalD
    for key in secondDict.keys():
        if type(secondDict[key]).__name__ == 'dict':
            plotTree(secondDict[key], cntrPt, str(key))
        else:
            plotTree.x0ff = plotTree.x0ff + 1.0 / plotTree.totalW
            plotNode(secondDict[key], (plotTree.x0ff, plotTree.y0ff), cntrPt, leafNode)
            plotMidText((plotTree.x0ff, plotTree.y0ff), cntrPt, str(key))
    plotTree.y0ff = plotTree.y0ff + 1.0 / plotTree.totalD


def createPlot(inTree):
    fig = plt.figure(1, facecolor='white')
    fig.clf()
    axprops = dict(xticks=[], yticks=[])
    createPlot.ax1 = plt.subplot(111, frameon=False, **axprops)
    plotTree.totalW = float(getNumLeafs(inTree))
    plotTree.totalD = float(getTreeDepth(inTree))
    plotTree.x0ff = -0.5 / plotTree.totalW
    plotTree.y0ff = 1.0
    plotTree(inTree, (0.5, 1.0), '')
    plt.show()


if __name__ == '__main__':
    df = pd.read_csv('../dataset/watermelon_4_3.csv')
    data = df.values[:, 1:].tolist()
    data_full = data[:]
    labels = df.columns.values[1:-1].tolist()
    labels_full = labels[:]
    myTree = createTree(data, labels, data_full, labels_full)
    createPlot(myTree)
